#!/usr/bin/env python3
"""
图片转LCD显示格式工具
功能：将图片转换为LCD显示所需的数据
- 水平扫描方式
- 支持RGB332(256色)或E6(4bpp)输出
- 可生成二进制数据或C语言数组
"""

import argparse
import os
from PIL import Image
import sys

EPD_BLACK = 0x00
EPD_WHITE = 0x01
EPD_YELLOW = 0x02
EPD_RED = 0x03
EPD_BLUE = 0x05
EPD_GREEN = 0x06
EPD_CLEAN = 0x07

# case '6-colors':
#     return [
#         { r: 0, g: 0, b: 0 },       // 榛戣壊 0x000000
#         { r: 255, g: 255, b: 255 }, // 鐧借壊 0xffffff
#         { r: 255, g: 255, b: 0 },   // 榛勮壊 0xffff00
#         { r: 255, g: 0, b: 0 },     // 绾㈣壊 0xff0000
#         { r: 0, g: 255, b: 0 },     // 缁胯壊 0x00ff00
#         { r: 0, g: 0, b: 255 }      // 钃濊壊 0x0000ff
#     ];

def rgb_to_rgb332(r, g, b):
    """
    将RGB888转换为RGB332格式(256色)
    R: 3位 (0-7)
    G: 3位 (0-7)
    B: 2位 (0-3)
    """
    r_3bit = (r >> 5) & 0x07  # 取高3位
    g_3bit = (g >> 5) & 0x07  # 取高3位
    b_2bit = (b >> 6) & 0x03  # 取高2位

    # 组合成8位: RRRGGGBB
    return (r_3bit << 5) | (g_3bit << 2) | b_2bit

def rgb_to_e6(r, g, b):
    """
    将RGB888转换为E6格式(6色)
    颜色映射:
        黑色: 0x00
        白色: 0x01
        黄色: 0x02
        红色: 0x03
        蓝色: 0x05
        绿色: 0x06
        清除: 0x07
    """
    if r == 0 and g == 0 and b == 0:
        return EPD_BLACK
    elif r == 255 and g == 255 and b == 255:
        return EPD_WHITE
    elif r == 255 and g == 255 and b == 0:
        return EPD_YELLOW
    elif r == 255 and g == 0 and b == 0:
        return EPD_RED
    elif r == 0 and g == 0 and b == 255:
        return EPD_BLUE
    elif r == 0 and g == 255 and b == 0:
        return EPD_GREEN
    else:
        return EPD_CLEAN

def image_to_lcd_array(
    image_path,
    output_path=None,
    array_name=None,
    max_width=None,
    max_height=None,
    rotate=0,
    output_format="rgb332_c",
):
    """
    将图片转换为LCD数组格式

    参数:
        image_path: 输入图片路径
        output_path: 输出文件路径，如果为None则输出到stdout
        array_name: C数组名称（仅在输出C数组时使用）
        max_width: 最大宽度，如果指定则调整图片大小
        max_height: 最大高度，如果指定则调整图片大小
        rotate: 旋转角度，可选值: 0, 90, 180, 270
        output_format: 输出格式，可选: rgb332, rgb332_c, e6, e6_c
    """
    fmt = output_format.lower()
    is_c_output = fmt.endswith("_c")
    color_format = fmt.replace("_c", "")
    if color_format not in ("rgb332", "e6"):
        print(f"错误: 不支持的输出格式 '{output_format}'", file=sys.stderr)
        return False

    if array_name is None:
        if output_path:
            base_name = os.path.basename(output_path)
            array_name = os.path.splitext(base_name)[0] + "_data"
        else:
            array_name = "image_data"
    try:
        # 打开图片
        img = Image.open(image_path)
        print(f"原始图片尺寸: {img.size[0]} x {img.size[1]}", file=sys.stderr)

        # 转换为RGB模式
        if img.mode != "RGB":
            img = img.convert("RGB")
            print(f"已转换为RGB模式", file=sys.stderr)

        # 旋转图片（如果需要）
        if rotate != 0:
            if rotate == 90:
                img = img.rotate(-90, expand=True)
                print(f"已旋转90度（顺时针）", file=sys.stderr)
            elif rotate == 180:
                img = img.rotate(180, expand=True)
                print(f"已旋转180度", file=sys.stderr)
            elif rotate == 270:
                img = img.rotate(90, expand=True)
                print(f"已旋转270度（顺时针）", file=sys.stderr)
            else:
                print(
                    f"警告: 不支持的旋转角度 {rotate}，仅支持 0, 90, 180, 270",
                    file=sys.stderr,
                )
            print(f"旋转后尺寸: {img.size[0]} x {img.size[1]}", file=sys.stderr)

        # 调整大小（如果需要）
        if max_width or max_height:
            # 计算缩放比例，保持宽高比
            ratio = 1.0
            if max_width and img.size[0] > max_width:
                ratio = min(ratio, max_width / img.size[0])
            if max_height and img.size[1] > max_height:
                ratio = min(ratio, max_height / img.size[1])

            if ratio < 1.0:
                new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
                img = img.resize(new_size, Image.LANCZOS)
                print(f"调整后尺寸: {img.size[0]} x {img.size[1]}", file=sys.stderr)

        width, height = img.size
        pixels = img.load()

        # 根据颜色格式生成数据
        data_bytes = []
        if color_format == "e6":
            color_desc = "E6 (4bpp，6色，两个像素1字节)"
            high_nibble = True
            current_byte = 0
            for y in range(height):
                for x in range(width):
                    r, g, b = pixels[x, y]
                    nibble = rgb_to_e6(r, g, b) & 0x0F
                    if high_nibble:
                        current_byte = nibble << 4
                        high_nibble = False
                    else:
                        current_byte |= nibble
                        data_bytes.append(current_byte)
                        high_nibble = True
                        current_byte = 0

            if not high_nibble:
                # Odd pixel count leaves low nibble empty; pad with 0
                data_bytes.append(current_byte)
        else:
            color_desc = "RGB332 (256色)"
            for y in range(height):
                for x in range(width):
                    r, g, b = pixels[x, y]
                    data_bytes.append(rgb_to_rgb332(r, g, b))

        byte_count = len(data_bytes)

        if not is_c_output:
            if not output_path:
                print("错误: 二进制输出需要使用 -o 指定文件", file=sys.stderr)
                return False
            with open(output_path, "wb") as f:
                f.write(bytes(data_bytes))
            print(f"已生成二进制数据文件: {output_path}", file=sys.stderr)
            print(f"颜色格式: {color_desc}", file=sys.stderr)
            print(f"数据大小: {byte_count} 字节", file=sys.stderr)
            return True

        # 生成C数组内容
        output_lines = []
        output_lines.append("/*")
        output_lines.append(f" * 图片文件: {os.path.basename(image_path)}")
        output_lines.append(f" * 图片尺寸: {width} x {height}")
        output_lines.append(f" * 颜色格式: {color_desc}")
        output_lines.append(f" * 扫描方式: 水平扫描")
        output_lines.append(f" * 数据大小: {byte_count} 字节")
        output_lines.append(" */")
        output_lines.append("")
        output_lines.append(f"#define {array_name.upper()}_WIDTH  {width}")
        output_lines.append(f"#define {array_name.upper()}_HEIGHT {height}")
        output_lines.append("")
        output_lines.append(f"const unsigned char {array_name}[{byte_count}] = {{")

        # 输出字节数据，每行最多16个字节
        line_buffer = "    "
        for index, byte in enumerate(data_bytes, start=1):
            line_buffer += f"0x{byte:02X}, "
            if index % 16 == 0:
                output_lines.append(line_buffer)
                line_buffer = "    "

        # 添加剩余的字节
        if line_buffer.strip() and line_buffer != "    ":
            line_buffer = line_buffer.rstrip(", ")
            output_lines.append(line_buffer)
        elif output_lines and output_lines[-1].startswith("    "):
            output_lines[-1] = output_lines[-1].rstrip(", ")

        output_lines.append("};")
        output_lines.append("")

        # 输出结果
        result = "\n".join(output_lines)

        if output_path:
            with open(output_path, "w", encoding="utf-8") as f:
                f.write(result)
            print(f"已生成C数组文件: {output_path}", file=sys.stderr)
            print(f"数组名称: {array_name}", file=sys.stderr)
            print(f"数据大小: {byte_count} 字节", file=sys.stderr)
        else:
            print(result)

        return True

    except FileNotFoundError:
        print(f"错误: 找不到文件 '{image_path}'", file=sys.stderr)
        return False
    except Exception as e:
        print(f"错误: {str(e)}", file=sys.stderr)
        return False
def main():
    parser = argparse.ArgumentParser(
        description="将图片转换为LCD显示所需的二进制数据或C数组（RGB332/E6，水平扫描）",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  %(prog)s input.jpg -o output.h
  %(prog)s photo.png -o image.h -n my_image
  %(prog)s logo.bmp -o logo.h -w 128 --height 64
  %(prog)s picture.jpg -r 90  # 顺时针旋转90度
  %(prog)s picture.jpg -f e6_c -o image_e6.h  # 输出E6 C数组
  %(prog)s picture.jpg -f rgb332 -o image_rgb332.bin  # 输出RGB332二进制
  %(prog)s picture.jpg  # 输出到stdout（默认C数组）
        """,
    )

    parser.add_argument("input", help="输入图片路径")
    parser.add_argument("-o", "--output", help="输出文件路径（默认输出到stdout）")
    parser.add_argument(
        "-n", "--name", default=None, help="C数组名称（默认: image_data）"
    )
    parser.add_argument("-w", "--width", type=int, help="最大宽度（保持宽高比）")
    parser.add_argument(
        "--height",
        type=int,
        dest="max_height",
        default=None,
        help="最大高度（保持宽高比）",
    )
    parser.add_argument(
        "-r",
        "--rotate",
        type=int,
        choices=[0, 90, 180, 270],
        default=0,
        help="旋转角度（顺时针），可选: 0, 90, 180, 270（默认: 0）",
    )
    parser.add_argument(
        "-f",
        "--format",
        choices=["rgb332", "rgb332_c", "e6", "e6_c"],
        default="rgb332_c",
        help="输出格式: rgb332/e6(二进制) 或 rgb332_c/e6_c(C数组)",
    )

    args = parser.parse_args()

    # 检查输入文件是否存在
    if not os.path.exists(args.input):
        print(f"错误: 文件不存在 '{args.input}'", file=sys.stderr)
        sys.exit(1)

    # 执行转换
    success = image_to_lcd_array(
        args.input,
        args.output,
        args.name,
        args.width,
        args.max_height,
        args.rotate,
        args.format,
    )

    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main()
